'''
Created on Feb 19, 2013

@author: nzhao
'''
import numpy as np

from traits.api import Instance, List
from math_qm.QMobj import QMobject

QUBIT_OPERATORS = ['I', 'X', 'Y', 'Z']

class QMoperator(QMobject):
    mat     = Instance(np.matrix)
    entries = List( QMobject )
    
    def __init__(self, nbody = 0, dim_list = [], mat = None, entries = [], qubit_operator = None):
        if mat is not None:
            if nbody == 0 or dim_list == []:
                print 'error: nbody and dim_list must be assigned!'
            else:
                QMobject.__init__(self, nbody, dim_list)
                self.mat = mat
                self.entries = []

        elif len(entries)>0:
            self.entries = entries
            nbody   = len(entries)
            dim_lst = [entry.dim for entry in entries]
            QMobject.__init__(self, nbody, dim_lst)
            
            prod_mat = 1
            for i in range(len(entries)):
                prod_mat = np.kron(prod_mat, entries[i].mat)
            self.mat = prod_mat
        
        elif qubit_operator in QUBIT_OPERATORS:
            QMobject.__init__(self, 1, [2])
            self.entries = []
            if qubit_operator is 'I':
                self.mat = np.matrix('1.0, 0.0; 0.0, 1.0')
            elif qubit_operator is 'X':
                self.mat = np.matrix('0.0, 0.5; 0.5, 0.0')
            elif qubit_operator is 'Y':
                self.mat = np.matrix('0.0, -1.j; 1.j, 0.0')
            elif qubit_operator is 'Z':
                self.mat = np.matrix('1., 0.0; 0.0, -1.0')
        
        else:
            print 'error: Unable to complete initialization.'